import cv2
import numpy as np
import os
import re

def merge_images():
    # 输入目录
    input_dir = "pic"
    
    # 确保输入目录存在
    if not os.path.exists(input_dir):
        print(f"Error: Directory '{input_dir}' does not exist")
        return
    
    # 获取所有图片文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    
    if not image_files:
        print(f"No image files found in {input_dir}")
        return
    
    # 自然排序图片文件名
    def natural_sort_key(s):
        return [int(text) if text.isdigit() else text.lower()
                for text in re.split('([0-9]+)', s)]
    
    image_files = sorted(image_files, key=natural_sort_key)
    print(f"Found {len(image_files)} images to merge")
    for img_file in image_files:
        print(img_file)
    
    # 读取第一张图片来获取宽度
    first_image = cv2.imread(os.path.join(input_dir, image_files[0]))
    if first_image is None:
        print(f"Error: Could not read first image {image_files[0]}")
        return
    
    width = first_image.shape[1]
    
    # 如果图片数量小于等于5，直接合并
    if len(image_files) <= 5:
        merge_image_group(image_files, input_dir, width, "merged_result.jpg")
    else:
        # 将图片分组，每组5张
        groups = [image_files[i:i+5] for i in range(0, len(image_files), 5)]
        
        # 处理每一组图片
        for i, group in enumerate(groups):
            output_filename = f"merged_result_{i+1:02d}.jpg"
            merge_image_group(group, input_dir, width, output_filename)
            print(f"Merged group {i+1} saved as {output_filename}")

def merge_image_group(image_files, input_dir, width, output_filename):
    """合并一组图片"""
    total_height = 0
    images = []
    
    # 读取所有图片并计算总高度
    for img_file in image_files:
        img_path = os.path.join(input_dir, img_file)
        img = cv2.imread(img_path)
        
        if img is None:
            print(f"Error: Could not read image {img_file}")
            continue
        
        # 确保所有图片宽度一致
        if img.shape[1] != width:
            img = cv2.resize(img, (width, int(img.shape[0] * width / img.shape[1])))
        
        images.append(img)
        total_height += img.shape[0]
    
    # 创建空白画布
    merged_image = np.zeros((total_height, width, 3), dtype=np.uint8)
    
    # 在画布上垂直拼接图片
    current_y = 0
    for img in images:
        h = img.shape[0]
        merged_image[current_y:current_y + h, :] = img
        current_y += h
    
    # 保存结果
    output_path = os.path.join(input_dir, output_filename)
    cv2.imwrite(output_path, merged_image)
    print(f"Merged image saved to {output_path}")

if __name__ == "__main__":
    merge_images() 